#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Elasticsearch内存管理器
扩展qwen_agent的memory模块，支持Elasticsearch文档存储和检索
"""

import os
import logging
from typing import List, Dict, Any, Optional, Union
from qwen_agent.memory.memory import Memory
from qwen_agent.tools.elasticsearch_retrieval import ElasticsearchRetrieval
from qwen_agent.tools.doc_parser import DocParser

logger = logging.getLogger(__name__)


class ElasticsearchMemory(Memory):
    """
    基于Elasticsearch的内存管理器
    替代默认的Memory实现，提供高性能文档存储和检索
    """
    
    def __init__(self, 
                 llm: Dict[str, Any],
                 files: Optional[List[str]] = None,
                 es_config: Optional[Dict[str, Any]] = None,
                 **kwargs):
        """
        初始化Elasticsearch内存管理器
        
        Args:
            llm: LLM配置
            files: 要索引的文件列表
            es_config: Elasticsearch配置
            **kwargs: 其他参数
        """
        # 不调用父类初始化，避免创建默认的retrieval工具
        self.llm = llm
        self.files = files or []
        self.es_config = es_config or {}
        
        # 设置默认ES配置
        default_es_config = {
            'es_host': 'https://localhost',
            'es_port': 9200,
            'es_user': 'elastic',
            'es_password': 'ncUg27UP5aMSzSElVXAX',
            'index_name': 'qwen_agent_docs',
            'verify_certs': False,
            'chunk_size': 500,
            'chunk_overlap': 50,
            'max_results': 5
        }
        
        # 合并配置
        self.es_config = {**default_es_config, **self.es_config}
        
        # 初始化组件
        self.doc_parser = DocParser(cfg=kwargs.get('cfg', {}))
        self.retrieval_tool = ElasticsearchRetrieval(cfg=self.es_config)
        
        # 索引文件
        if self.files:
            self._index_files()
        
        logger.info("✅ ElasticsearchMemory初始化完成")
    
    def _index_files(self):
        """索引文件"""
        try:
            logger.info(f"📚 开始索引 {len(self.files)} 个文件")
            success = self.retrieval_tool.index_files(self.files)
            
            if success:
                logger.info("✅ 文件索引完成")
            else:
                logger.error("❌ 文件索引失败")
                
        except Exception as e:
            logger.error(f"❌ 索引文件时出错: {str(e)}")
    
    def get_tool_list(self) -> List[Any]:
        """获取工具列表"""
        return [self.retrieval_tool]
    
    def retrieve(self, query: str, **kwargs) -> List[Dict[str, Any]]:
        """
        检索相关文档
        
        Args:
            query: 查询字符串
            **kwargs: 其他参数
            
        Returns:
            检索结果列表
        """
        try:
            k = kwargs.get('k', self.es_config.get('max_results', 5))
            results = self.retrieval_tool.search(query, k=k)
            
            logger.info(f"🔍 检索查询: {query}")
            logger.info(f"📊 找到 {len(results)} 个相关文档")
            
            return results
            
        except Exception as e:
            logger.error(f"❌ 检索失败: {str(e)}")
            return []
    
    def add_file(self, file_path: str) -> bool:
        """
        添加新文件到索引
        
        Args:
            file_path: 文件路径
            
        Returns:
            是否成功
        """
        try:
            if file_path not in self.files:
                self.files.append(file_path)
            
            success = self.retrieval_tool.index_files([file_path])
            
            if success:
                logger.info(f"✅ 成功添加文件: {os.path.basename(file_path)}")
            else:
                logger.error(f"❌ 添加文件失败: {file_path}")
            
            return success
            
        except Exception as e:
            logger.error(f"❌ 添加文件时出错: {str(e)}")
            return False
    
    def add_files(self, file_paths: List[str]) -> bool:
        """
        批量添加文件到索引
        
        Args:
            file_paths: 文件路径列表
            
        Returns:
            是否成功
        """
        try:
            # 添加到文件列表
            for file_path in file_paths:
                if file_path not in self.files:
                    self.files.append(file_path)
            
            # 批量索引
            success = self.retrieval_tool.index_files(file_paths)
            
            if success:
                logger.info(f"✅ 成功添加 {len(file_paths)} 个文件")
            else:
                logger.error(f"❌ 批量添加文件失败")
            
            return success
            
        except Exception as e:
            logger.error(f"❌ 批量添加文件时出错: {str(e)}")
            return False
    
    def search_documents(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        """
        搜索文档（兼容性方法）
        
        Args:
            query: 查询字符串
            k: 返回结果数量
            
        Returns:
            搜索结果
        """
        return self.retrieve(query, k=k)
    
    def get_index_stats(self) -> Dict[str, Any]:
        """获取索引统计信息"""
        try:
            if not self.retrieval_tool.es_client:
                return {"error": "ES客户端未连接"}
            
            stats = self.retrieval_tool.es_client.indices.stats(
                index=self.retrieval_tool.index_name
            )
            
            index_stats = stats['indices'].get(self.retrieval_tool.index_name, {})
            
            return {
                "index_name": self.retrieval_tool.index_name,
                "document_count": index_stats.get('total', {}).get('docs', {}).get('count', 0),
                "store_size": index_stats.get('total', {}).get('store', {}).get('size_in_bytes', 0),
                "indexed_files": len(self.files)
            }
            
        except Exception as e:
            logger.error(f"❌ 获取索引统计失败: {str(e)}")
            return {"error": str(e)}
    
    def clear_index(self) -> bool:
        """清空索引"""
        try:
            if not self.retrieval_tool.es_client:
                return False
            
            # 删除索引
            if self.retrieval_tool.es_client.indices.exists(index=self.retrieval_tool.index_name):
                self.retrieval_tool.es_client.indices.delete(index=self.retrieval_tool.index_name)
                logger.info(f"🗑️ 删除索引: {self.retrieval_tool.index_name}")
            
            # 重新创建索引
            self.retrieval_tool._create_index_if_not_exists()
            
            # 清空文件列表
            self.files = []
            
            logger.info("✅ 索引清空完成")
            return True
            
        except Exception as e:
            logger.error(f"❌ 清空索引失败: {str(e)}")
            return False


def create_elasticsearch_memory(llm: Dict[str, Any],
                              files: Optional[List[str]] = None,
                              es_config: Optional[Dict[str, Any]] = None,
                              **kwargs) -> ElasticsearchMemory:
    """
    创建Elasticsearch内存管理器
    
    Args:
        llm: LLM配置
        files: 文件列表
        es_config: ES配置
        **kwargs: 其他参数
        
    Returns:
        ElasticsearchMemory实例
    """
    return ElasticsearchMemory(
        llm=llm,
        files=files,
        es_config=es_config,
        **kwargs
    )